import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from mpl_toolkits.mplot3d import Axes3D  


def sigmoid_lm1_r0(x, c=0, k=5):  
    return -1 / (1 + np.exp(k * (x - c)))  

#根据流函数计算u,v
def psi_to_vel(psi,u,v,dx,dy):
    #内点
    u[1:-1,1:-1]=(psi[2:, 1:-1] - psi[:-2, 1:-1])/2/dy
    v[1:-1,1:-1]=-(psi[1:-1, 2:] - psi[1:-1, :-2])/2/dx
    #边界
    u[:,0]=u[:,-2]
    u[:,-1]=u[:,1]
    u[-1, 1:-1] = u_up
    u[0, 1:-1] = 0

    v[:,0]=v[:,-2]
    v[:,-1]=v[:,1]
    v[-1, 1:-1] = 0 
    v[0, 1:-1] = v[1, 1:-1] 
    return u,v

#设置参数
Nx = 4+2
Ny = 4+2
Lx = Ly = 1
u_up=0
x = np.linspace(0, Lx, Nx)  #linspace生成的点列是包含头和尾的。
dx = x[-1] - x[-2]  #网格长
y = np.linspace(0, Ly, Ny)
dy = y[-1] - y[-2]
nu = 0.01
dt = 0.005  #时间步长
ntime=1000  #时间步数
ndiag=20 #每隔多少时间输出一次
Re=1/nu
xx, yy = np.meshgrid(x, y)  

#初始条件设置
#设置初始流函数场
psi=np.zeros_like(yy)
phi0_y=sigmoid_lm1_r0(y,c=Ly/2,k=20/Ly)
for i in range(0,Nx):
    for j in range(0,Ny):
        psi[j][i]=phi0_y[j]

print(psi)

#设置初始速度场
u = np.zeros_like(xx)
v = np.zeros_like(yy)
u,v=psi_to_vel(psi,u,v,dx,dy)
u[-1, 1:-1] = u_up    #最上方速度1。这里设置速度时排除了首尾格点。numpy可以一个数组等于一个数，这样把这个数赋值给这个数组

#绘制截面速度图
flag_u_ysection=True
if flag_u_ysection:
    #增设画图部分。
    plt.plot(y,u[:,Nx//2]) 
    plt.show()

#绘制流函数三维图
flag_axe3D=True
if flag_axe3D:
    # 创建三维图形和坐标轴  
    fig = plt.figure()  
    ax = fig.add_subplot(111, projection='3d')  
    
    # 绘制三维曲面  
    ax.plot_surface(xx, yy, psi, cmap='viridis', linewidth=0)  # cmap参数用于设置颜色映射  
    
    # 添加标题和坐标轴标签  
    ax.set_title('3D Plot of psi')  
    ax.set_xlabel('X')  
    ax.set_ylabel('Y')  
    ax.set_zlabel('Psi')  
    
    # 显示图形  
    plt.show()

flag_quiver=True
if flag_quiver:
    # 创建一个新的figure  
    plt.figure()  
    
    # 使用quiver函数绘制矢量图  
    Q = plt.quiver(xx, yy, u, v, units='width', scale=10, color='b')  
    
    # 添加标题和坐标轴标签  
    plt.title('2D Vector Field')  
    plt.xlabel('X')  
    plt.ylabel('Y')  
    
    # 显示图形  
    plt.show()